package datastructure.trie;

import java.util.TreeMap;

public class Trie {

    class Node {
        private boolean isWord;
        private TreeMap<Character, Node> next;

        public Node(boolean isWord) {
            this.isWord = isWord;
            next = new TreeMap<>();
        }

        public Node() {
            this(false);
        }
    }

    // 定义跟节点
    private Node root;
    private int size;

    public Trie() {
        root = new Node();
        size = 0;
    }

    public int getSize() {
        return size;
    }

    // 添加元素单词 ;使用循环
    public void add(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (cur.next.get(c) == null) {
                cur.next.put(c, new Node());
            }
            cur = cur.next.get(c);
        }
        // 如果执行到这里,说明cur指向的是最后一个单词的节点
        // 判断当前节点原本是否已经是一个单词节点
        if (!cur.isWord) {
            cur.isWord = true;
            // 维护大小
            size++;
        }
    }

    // 查找一个元素是否已经存在在该集合中
    public boolean contains(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            // 如果第一个单词都不能匹配,说明单词不存在,直接返回false
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }
        // 执行到这里,我们就找到了cur指向这个单词的最后一个位置
        return cur.isWord;// 应该返回的是当前节点是否是一个单词的节点
    }

    // 查询是否在Trie中有单词以prefix为前缀
    public boolean isPrefix(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);

            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }

        return true;
    }
}
